package com.dingding.aop;

import cn.hutool.core.util.StrUtil;
import cn.hutool.http.HtmlUtil;
import cn.hutool.json.JSONUtil;
import lombok.extern.slf4j.Slf4j;

import javax.servlet.ReadListener;
import javax.servlet.ServletInputStream;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletRequestWrapper;
import java.io.*;
import java.nio.charset.Charset;
import java.util.HashMap;
import java.util.LinkedHashMap;
import java.util.Map;

/**
 * 构建可重复读取inputStream的request
 *
 * @author harbour
 */
@Slf4j
public class XssHttpServletRequestWrapper extends HttpServletRequestWrapper {

  private final byte[] body;

  public XssHttpServletRequestWrapper(HttpServletRequest request) throws IOException {
    super(request);
    body = inputStream2String(request.getInputStream()).getBytes(Charset.forName("UTF-8"));
  }

  @Override
  public String getParameter(String name) {
    String value = super.getParameter(name);
    if (!StrUtil.hasEmpty(value)) {
      value = HtmlUtil.filter(value);
    }
    return value;
  }

  @Override
  public String[] getParameterValues(String name) {
    String[] values = super.getParameterValues(name);
    if (values != null) {
      for (int i = 0; i < values.length; i++) {
        String value = values[i];
        if (!StrUtil.hasEmpty(value)) {
          value = HtmlUtil.filter(value);
        }
        values[i] = value;
      }
    }
    return values;
  }

  @Override
  public Map<String, String[]> getParameterMap() {
    Map<String, String[]> parameters = super.getParameterMap();
    Map<String, String[]> map = new LinkedHashMap<>();
    if (parameters != null) {
      for (String key : parameters.keySet()) {
        String[] values = parameters.get(key);
        for (int i = 0; i < values.length; i++) {
          String value = values[i];
          if (!StrUtil.hasEmpty(value)) {
            value = HtmlUtil.filter(value);
          }
          values[i] = value;
        }
        map.put(key, values);
      }
    }
    return map;
  }

  @Override
  public String getHeader(String name) {
    String value = super.getHeader(name);
    if (!StrUtil.hasEmpty(value)) {
      value = HtmlUtil.filter(value);
    }
    return value;
  }

  @Override
  public BufferedReader getReader() throws IOException {
    return new BufferedReader(new InputStreamReader(getInputStream()));
  }

  @Override
  public ServletInputStream getInputStream() throws IOException {

    Map<String, Object> map = JSONUtil.parseObj(new String(body));
    Map<String, Object> resultMap = new HashMap(map.size());
    for (String key : map.keySet()) {
      Object val = map.get(key);
      if (map.get(key) instanceof String) {
        resultMap.put(key, HtmlUtil.filter(val.toString()));
      } else {
        resultMap.put(key, val);
      }
    }
    String str = JSONUtil.toJsonStr(resultMap);
    try (final ByteArrayInputStream bain = new ByteArrayInputStream(str.getBytes()); ) {

      return new ServletInputStream() {
        @Override
        public int read() throws IOException {
          return bain.read();
        }
        @Override
        public boolean isFinished() {
          return false;
        }

        @Override
        public boolean isReady() {
          return false;
        }

        @Override
        public void setReadListener(ReadListener readListener) {}
      };
    }
  }

  /**
   * 将 inputStream 里的数据读取出来并转换成字符串
   *
   * @param inputStream inputStream
   * @return String
   */
  public String inputStream2String(InputStream inputStream) {
    StringBuilder sb = new StringBuilder();
    BufferedReader reader = null;
    try {
      reader = new BufferedReader(new InputStreamReader(inputStream, Charset.forName("UTF-8")));
      String line;
      while ((line = reader.readLine()) != null) {
        sb.append(line);
      }
    } catch (IOException e) {
      sb.append("get body params fail");
      log.error(e.getMessage());
    } finally {
      if (reader != null) {
        try {
          reader.close();
        } catch (IOException e) {
          log.error(e.getMessage());
        }
      }
    }
    return sb.toString();
  }


}
